# security/token_service.py
import os
from dotenv import load_dotenv
from datetime import datetime, timedelta
from typing import Dict, Optional, Tuple
from jose import jwk, jwt
import json
from fastapi import HTTPException, status
from jwt import PyJWTError


# 环境变量配置（实际项目中应使用 .env 或配置中心）
class TokenConfig:
	SECRET_KEY = os.getenv("JWT_SECRET", "your_strong_secret_key")
	JWT_PUBLIC_KEY = os.getenv("JWT_PUBLIC_KEY")
	ALGORITHM = os.getenv("JWT_ALGORITHM", "RS256")
	ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE", 12*60))
	REFRESH_TOKEN_EXPIRE_DAYS = int(os.getenv("REFRESH_TOKEN_EXPIRE", 7))
	TOKEN_REFRESH_LEEWAY = int(os.getenv("TOKEN_REFRESH_LEEWAY", 5))  # 提前刷新时间（分钟）

JWT_PUBLIC_KEY = os.getenv("JWT_PUBLIC_KEY")
# 将字符串转换为JSON格式
PUBLIC_KEY = json.loads(JWT_PUBLIC_KEY)

class TokenService:
	@staticmethod
	def create_tokens(user_id: str, payload: Optional[Dict] = None) -> Tuple[str, str]:
		"""
		创建访问令牌和刷新令牌对

		:param user_id: 用户唯一标识
		:param payload: 附加的JWT负载数据
		:return: (access_token, refresh_token)
		"""
		# 基础负载数据
		base_payload = {
			"sub": user_id,
			"iat": datetime.utcnow()
		}
		
		# 合并自定义负载
		token_payload = {**base_payload, **(payload or {})}
		
		# 创建访问令牌
		access_payload = token_payload.copy()
		access_expire = datetime.utcnow() + timedelta(
			minutes=TokenConfig.ACCESS_TOKEN_EXPIRE_MINUTES
		)
		access_payload.update({
			"exp": access_expire,
			"type": "access"
		})
		access_token = jwt.encode(
			access_payload,
			TokenConfig.SECRET_KEY,
			algorithm=TokenConfig.ALGORITHM
		)
		
		# 创建刷新令牌
		refresh_payload = token_payload.copy()
		refresh_expire = datetime.utcnow() + timedelta(
			days=TokenConfig.REFRESH_TOKEN_EXPIRE_DAYS
		)
		refresh_payload.update({
			"exp": refresh_expire,
			"type": "refresh"
		})
		refresh_token = jwt.encode(
			refresh_payload,
			TokenConfig.SECRET_KEY,
			algorithm=TokenConfig.ALGORITHM
		)
		
		return access_token, refresh_token
	
	@staticmethod
	def verify_token(token: str) -> dict:
		"""
		验证JWT令牌并返回负载

		:param token: JWT令牌字符串
		:return: 解码后的令牌负载
		:raises HTTPException: 令牌无效或过期
		"""
		try:
			# 解码token但不验证，以获取header信息
			payload = jwt.decode(
				token,
				# 使用IdentityServer4 公钥进行验证
				PUBLIC_KEY,
				algorithms=[TokenConfig.ALGORITHM],
				options={"verify_exp": True}  # 显式启用过期验证
			)
			return payload
		except jwt.ExpiredSignatureError:
			raise HTTPException(
				status_code=status.HTTP_401_UNAUTHORIZED,
				detail="Token has expired",
				headers={"WWW-Authenticate": "Bearer"},
			)
		except (jwt.InvalidTokenError, PyJWTError):
			raise HTTPException(
				status_code=status.HTTP_401_UNAUTHORIZED,
				detail="Could not validate credentials",
				headers={"WWW-Authenticate": "Bearer"},
			)
	
	@staticmethod
	def refresh_access_token(refresh_token: str) -> str:
		"""
		使用刷新令牌生成新的访问令牌

		:param refresh_token: 有效的刷新令牌
		:return: 新的访问令牌
		:raises HTTPException: 刷新令牌无效或过期
		"""
		try:
			# 验证刷新令牌
			payload = TokenService.verify_token(refresh_token)
			
			# 检查令牌类型
			if payload.get("type") != "refresh":
				raise HTTPException(
					status_code=status.HTTP_401_UNAUTHORIZED,
					detail="Invalid token type"
				)
			
			# 创建新访问令牌（保留原始负载数据）
			new_payload = {k: v for k, v in payload.items()
			               if k not in ["exp", "iat", "type"]}
			access_token, _ = TokenService.create_tokens(
				user_id=payload["sub"],
				payload=new_payload
			)
			
			return access_token
		except HTTPException as e:
			raise e
		except Exception as e:
			raise HTTPException(
				status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
				detail=f"Token refresh failed: {str(e)}"
			)
	
	@staticmethod
	def should_refresh_token(token: str) -> bool:
		"""
		检查访问令牌是否需要刷新（基于时间余量）

		:param token: 访问令牌
		:return: 是否需要刷新
		"""
		try:
			payload = jwt.decode(
				token,
				TokenConfig.SECRET_KEY,
				algorithms=[TokenConfig.ALGORITHM],
				options={"verify_exp": False}  # 不验证过期
			)
			
			# 检查过期时间
			exp_timestamp = payload.get("exp")
			if not exp_timestamp:
				return False
			
			expire_time = datetime.utcfromtimestamp(exp_timestamp)
			time_remaining = expire_time - datetime.utcnow()
			
			# 如果剩余时间小于刷新阈值，则需要刷新
			return time_remaining < timedelta(minutes=TokenConfig.TOKEN_REFRESH_LEEWAY)
		except (jwt.InvalidTokenError, PyJWTError):
			return False